from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from configs import CHATLAW_TEXT2VEC_MODEL_PATH, device
from tqdm import trange
import numpy as np
import torch
import logging
from torch import Tensor

logger = logging.getLogger(__name__)

show_progress_bar = (logger.getEffectiveLevel()==logging.INFO or logger.getEffectiveLevel()==logging.DEBUG)
model = SentenceTransformer(CHATLAW_TEXT2VEC_MODEL_PATH)
model.eval()
model = model.to(device)


def batch_to_device(batch):
    for key in batch:
        if isinstance(batch[key], Tensor):
            batch[key] = batch[key].to(device)
    return batch


def chatlaw_text2vec_encode(text_list, is_normalize=True):
    if not isinstance(text_list, list):
        text_list = [text_list]
        flag = 0
    else:
        flag = 1
    batch_size = 32
    length_sorted_idx = np.argsort([-len(sen) for sen in text_list])
    sentences_sorted = [text_list[idx] for idx in length_sorted_idx]
    all_embeddings = []
    with torch.no_grad():
        for start_index in trange(0, len(text_list), batch_size, desc="Batches", disable=not show_progress_bar):
            sentences_batch = sentences_sorted[start_index:start_index+batch_size]
            encoded_input = model.tokenize(sentences_batch)
            encoded_input = batch_to_device(encoded_input)
            fea = model(encoded_input)
            embeddings = fea["sentence_embedding"]
            embeddings = embeddings.detach()
            if is_normalize:
                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            embeddings = embeddings.cpu()
            all_embeddings.extend(embeddings)
    all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]
    all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
    if flag == 0:
        return all_embeddings[0]
    else:
        return all_embeddings
    return all_embeddings















# from sentence_transformers import SentenceTransformer,  LoggingHandler, losses, models, util
# from sentence_transformers.util import cos_sim
# from history_qa.configs import CHATLAW_TEXT2VEC_MODEL_PATH
# import numpy as np
# import torch
# # CHATLAW_TEXT2VEC_MODEL_PATH = "../embedding_models/ChatLaw-Text2Vec"

# device = torch.device("cuda:3")
# chatlaw_text2vec_model = SentenceTransformer(CHATLAW_TEXT2VEC_MODEL_PATH).to(device)


# if __name__ == '__main__':
#     # sentences = "昆仑健康爱多保重大疾病保险 释义  （十）轻症疾病所属组别： 第一组 第二组 第三组 第四组 1.极早期恶性肿瘤或恶性病变 2.慢性肾功能损害-肾功能衰竭期 3.肾脏切除 4.病毒性肝炎导致的肝硬化 5.慢性肝功能衰竭 6.糖尿病导致单足截除 7.中度溃疡性结肠炎 8.肺功能衰竭 9.肝叶切除 10.单侧肺脏切除 11.双侧卵巢或睾丸切除术 12.因肾上腺皮脂腺瘤切除肾上腺 13.胆道重建手术 14.可逆性再生障碍性贫血 15.中度严重克隆症 1.轻微脑中风后遗症 2.脑垂体瘤、脑囊肿、脑动脉瘤及脑血管瘤 3.中度瘫痪 4.中度阿尔茨海默病 5.中度脑损伤 6.中度帕金森氏病 7.中度运动神经元病 8.中度严重脊髓灰质炎 9.中度肌营养不良症 10.中度脑炎或脑膜炎后遗症 11.植入大脑内分流器 12.微创颅脑手术 13.深度昏迷 72 小时 14.重度头部外伤 15.严重阻塞性睡眠窒息症 1.不典型的急性心肌梗塞 2.微创冠状动脉搭桥术 3.冠状动脉介入手术（非开胸手术） 4.激光心肌血运重建术 5.心脏瓣膜介入手术（非开胸手术） 6.原发性肺动脉高压 7.主动脉内手术（非开胸手术） 8.早期象皮病 9.心包膜切除术 10.植入心脏起搏器 11.早期原发性心肌病 12.腔静脉过滤器植入术 13.特定周围动脉疾病的血管介入治疗 1.单个肢体缺失 2.单耳失聪 3.人工耳蜗植入术 4.听力严重受损 5.视力严重受损 6.角膜移植 7.单目失明 8.全身较小面积Ⅲ度烧伤 9.严重甲型及乙型血友病 10.中度类风湿性关节炎 11.早期系统性硬皮病 12.轻度面部烧伤 13.面部重建手术 （十一）全残：本合同所定义的全残是指至少满足下列情形之一者： 1.双目永久完全（注①）失明（注②）； 2.两上肢腕关节以上或两下肢踝关节以上缺失； 3.一上肢腕关节以上及一下肢踝关节以上缺失； 4.一目永久完全失明及一上肢腕关节以上缺失； 5.一目永久完全失明及一下肢踝关节以上缺失； 6.四肢关节机能永久完全丧失（注③）； 7.咀嚼、吞咽机能永久完全丧失（注④）； 8.中枢神经系统机能或胸、腹部脏器机能极度障碍，终身不能从事任何工作，为维持生命必要的日常生活活动，全需他人扶助（注⑤）。"
#     sentence1 = "合同法(1999-03-15): 第二百零六条 借款人应当按照约定的期限返还借款。对借款期限没有约定或者约定不明确，依照本法第六十一条的规定仍不能确定的，借款人可以随时返还；贷款人可以催告借款人在合理期限内返还。"
#     sentence2 = "请问如果借款没还怎么办。"
#     encoded_sentence1 = chatlaw_text2vec_model.encode(sentence1)
#     encoded_sentence2 = chatlaw_text2vec_model.encode(sentence2)
#     # print(cos_sim(encoded_sentence1, encoded_sentence2))
#     v = chatlaw_text2vec_model.encode([sentence1, sentence2])
#     print(np.linalg.norm(v, axis=1, keepdims=True).shape)
#     print(v / np.linalg.norm(v, axis=1, keepdims=True))
    
